"""
model name : 深度学习
file       : model.py
information:
    author : OuYang
    time   : 2025/1/13
"""

import torch
from torch import nn
import torchvision

# Cancel waring
import warnings

warnings.filterwarnings("ignore", category=UserWarning)


class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False, dropout_p=0.5):
        super(Conv, self).__init__()

        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=bias
        )

        self.bn = nn.BatchNorm2d(out_channels)

        self.act = nn.LeakyReLU()

        self.dropout = nn.Dropout2d(dropout_p)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        x = self.dropout(x)
        return x


# YOLO network
class YOLO(nn.Module):
    def __init__(self,
                 in_channels=3,
                 num_classes=3,
                 backbone='resnet18',
                 imgsz=448,
                 s=7,
                 b=2,
                 pretrained=True
                 ):
        super(YOLO, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.imgsz = imgsz
        self.s = s
        self.b = b
        self.pretrained = pretrained

        # Backbone
        self.backbone, out_features = self._build_backbone(backbone)
        for param in self.parameters():
            param.requires_grad = False

        # Yolo Back 4 layer
        self.conv_layer = nn.Sequential(
            nn.Conv2d(
                in_channels=out_features,
                out_channels=1024,
                kernel_size=3,
                padding=1
            ),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(),
            # nn.Dropout2d(0.5, inplace=True),
            nn.Conv2d(1024, 1024, 3, 2, 1),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(),
            # nn.Dropout2d(0.5, inplace=True),
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(),
            # nn.Dropout2d(0.5, inplace=True)
        )
        self.conn_layers = nn.Sequential(
            nn.Linear(self.s * self.s * 1024, 4096),
            nn.BatchNorm1d(4096),
            nn.LeakyReLU(),
            # nn.Dropout(0.5, inplace=True),
            nn.Linear(4096, self.s * self.s * self._cal_one_grid_cell_predictions()),
            nn.BatchNorm1d(self.s * self.s * self._cal_one_grid_cell_predictions()),
            nn.Sigmoid()
        )

        # Flatten
        self.flatten = nn.Flatten()

    def forward(self, x):
        x = self.backbone(x)
        x = self.conv_layer(x)
        x = self.flatten(x)
        x = self.conn_layers(x)
        x = torch.reshape(x, (-1, self.s, self.s, self._cal_one_grid_cell_predictions()))

        return x

    def _build_backbone(self, backbone_name):
        ret_backbone = None
        ret_out_features = None
        if backbone_name == 'resnet18':
            ret_backbone = torchvision.models.resnet18(pretrained=self.pretrained)
            ret_out_features = ret_backbone.fc.in_features
        elif backbone_name == 'resnet34':
            ret_backbone = torchvision.models.resnet34(pretrained=self.pretrained)
            ret_out_features = ret_backbone.fc.in_features
        elif backbone_name == 'resnet50':
            ret_backbone = torchvision.models.resnet50(pretrained=self.pretrained)
            ret_out_features = ret_backbone.fc.in_features
        elif backbone_name == 'resnet101':
            ret_backbone = torchvision.models.resnet101(pretrained=self.pretrained)
            ret_out_features = ret_backbone.fc.in_features
        elif backbone_name == 'resnet152':
            ret_backbone = torchvision.models.resnet152(pretrained=self.pretrained)
            ret_out_features = ret_backbone.fc.in_features

        # Delete Back 2 layer
        ret_backbone = nn.Sequential(
            *list(ret_backbone.children())[:-2]
        )

        return ret_backbone, ret_out_features

    def _cal_one_grid_cell_predictions(self):
        return self.b * 5 + self.num_classes


if __name__ == '__main__':
    model = YOLO(backbone='resnet34')

    # print(model)
    for name, param in model.named_parameters():
        print(name, param.requires_grad)

    # Test Data
    inputs = torch.randn(1, 3, 448, 448)
    outputs = model(inputs)
    print(outputs.shape)
